package com.fq.demo.controller;

import com.fq.demo.config.AiConfig;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.service.TokenStream;
import lombok.val;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDate;

/**
 * @Author 超chao
 * @Description TODO
 * @Date 2025/3/21/周五 11:23
 * @Version 1.0
 */

@RestController
@RequestMapping("/ai")
public class ChatController {
    @Autowired
    private QwenChatModel qwenChatModel;

    @Autowired
    private QwenStreamingChatModel qwenStreamingChatModel;

    @Autowired
    private AiConfig.Assistant assistant;

    @Autowired
    private AiConfig.AssistantUnique assistantUnique;
    @RequestMapping("/chat")
    public String chat(@RequestParam(defaultValue = "你好！") String message){
        return qwenChatModel.chat(message);
    }

    @RequestMapping(value = "/stream-chat",produces = "text/stream;charset=UTF-8")
    public Flux<String> stream(@RequestParam(defaultValue = "你好！") String message){
        Flux<String> objectFlux = Flux.create(fluxSink -> {
            qwenStreamingChatModel.chat(message, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    fluxSink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    fluxSink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    fluxSink.error(throwable);
                }
            });
        });
        return objectFlux;
    }

    /**
     * 记忆对话
     * @param message
     * @return
     */
    @RequestMapping("/memory-chat")
    public String memoryChat(@RequestParam(defaultValue = "你好！") String message){
        return assistant.chat(message);
    }

    /**
     * 记忆对话流式
     * @param message
     * @return
     */
    @RequestMapping(value = "/memory-stream-chat",produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryStream(@RequestParam(defaultValue = "你好！") String message){
        TokenStream stream = assistant.chatStream(message);
        return Flux.create(fluxSink -> {
            stream.onPartialResponse(fluxSink::next)
                    .onCompleteResponse(chatResponse -> fluxSink.complete())
                    .onError(fluxSink::error)
                    .start();
        });
    }

    /**
     * 记忆对话隔离
     * @param message
     * @param userId
     * @return
     */
    @RequestMapping(value = "/memoryId-stream-chat",produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryIdStream(@RequestParam(defaultValue = "你好！") String message,@RequestParam(defaultValue = "1") Integer userId){
        TokenStream stream = assistantUnique.chatStream(userId,message);
        return Flux.create(fluxSink -> {
            stream.onPartialResponse(fluxSink::next)
                    .onCompleteResponse(chatResponse -> fluxSink.complete())
                    .onError(fluxSink::error)
                    .start();
        });
    }

    /**
     * function-call 调用方法功能
     * @param message
     * @return
     */
    @RequestMapping("/memory-chat-tools")
    public String memoryChatTools(@RequestParam(defaultValue = "你好！") String message){
        return assistant.chat(message);
    }

    /**
     * 记忆对话流式
     * @param message
     * @return
     */
    @RequestMapping(value = "/memory-stream-chat-tools",produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryStreamTools(@RequestParam(defaultValue = "你好！") String message){
        TokenStream stream = assistant.chatStream(message, LocalDate.now().toString());
        return Flux.create(fluxSink -> {
            stream.onPartialResponse(fluxSink::next)
                    .onCompleteResponse(chatResponse -> fluxSink.complete())
                    .onError(fluxSink::error)
                    .start();
        });
    }
}
